fix(cambricon): stop index autotune from killing the process on untileable expand_shape - #5510
Open
tengqm wants to merge 1 commit into
Open
fix(cambricon): stop index autotune from killing the process on untileable expand_shape#5510tengqm wants to merge 1 commit into
tengqm wants to merge 1 commit into
Conversation
…eable expand_shape
The MLU backend's AutoTileForTritonPass raises a plain RuntimeError ('PassManager::run failed') on the tensor.expand_shape generated for BLOCK_SIZE0=4 x BLOCK_SIZE1=4096. Since RuntimeError is outside triton's autotuner catch list (OutOfResources / CompileTimeAssertionFailure / PTXASError), the exception propagated out of the tuner and crashed the whole process (EngineCore dead in vLLM).
Two-part fix, E2E-verified on MLU590 with flag_gems 5.3.4 + vLLM 0.20.2:
1) libentry bench(): any RuntimeError from _bench is treated as a non-candidate (inf) so tuning survives backend compiler failures on any config.
2) _cambricon index tune config: drop BLOCK_SIZE1=4096. With default num_stages=2 a 4096-wide tile exceeds NRAM (always OutOfResources) and BLOCK_SIZE0=4 x 4096 is the exact config that trips the compiler bug.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
E2E on cambricon MLU590 (vLLM 0.20.2 + flag_gems 5.3.4, inductor-wrapped context), the
indexop autotune kills the whole process instead of skipping a bad config:Root cause: the MLU backend's
AutoTileForTritonPassraises a plainRuntimeErroron thetensor.expand_shapegenerated for theBLOCK_SIZE0=4 x BLOCK_SIZE1=4096config (collapse[1, 4, 4096]).RuntimeErroris outside triton's autotuner catch list (OutOfResources,CompileTimeAssertionFailure,PTXASError— allTritonErrorsubclasses), so the exception escapes the tuner, propagates through libentry, and kills the caller (vLLMEngineCoredead).Standalone (non-wrapped) the config compiles fine; the failure only manifests in inductor-wrapped contexts where shapes are dynamic. Reproducer:
indexon a (15, 4096) bf16 tensor under inductor wrap,BLOCK_SIZE0=4,BLOCK_SIZE1=4096,num_stages=2.Fix (two-part, E2E-verified)
libentry.pybench(): anyRuntimeErrorfrom_benchis treated as a non-candidate (inf) so tuning survives backend compiler failures on any config. This is the general robustness fix — the failure mode is 'backend raises non-TritonError on an untileable config'._cambricon/tune_configs.yamlindex block: dropBLOCK_SIZE1=4096. With the defaultnum_stages=2, a 4096-wide tile always exceeds NRAM (standalone-tested:OutOfResources, can never win), andBLOCK_SIZE0=4 x 4096is the exact config that trips the compiler bug.Verified on MLU590: with the previous blacklist removed, the request that used to crash now returns correct output, the engine stays alive across repeated requests, and no errors appear in the log.
Note
The underlying
AutoTileForTritonPassexpand_shape failure is a vendor (cambricon triton-MLU) compiler bug; this change makes flag_gems resilient to it. If the vendor fixes it,BLOCK_SIZE1=4096can be re-added.closes: #5509
This PR was written in part with the assistance of generative AI.